import torch
import torch.nn as nn
import torch.nn.functional as F


class ConvBlock(nn.Module):
    def __init__(self, channel):
        super(ConvBlock, self).__init__()
        
        self.conv1 = nn.Conv2d(channel, 1, 1, 1, 0, bias=False)
        self.conv2 = nn.Conv2d(1, 1, 1, 1, 0, bias=False)
        
        for param in self.conv1.parameters():
            param.requires_grad = False
            
        for param in self.conv2.parameters():
            param.requires_grad = False
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        
        return x


class Perceptual_loss(nn.Module):
    def __init__(self, network=None, channel=3):
        super(Perceptual_loss, self).__init__()

        # pretrained network
        if network is not None:
            self.net = network
        else:
            self.net = ConvBlock(channel)
            
        self.loss = nn.L1Loss()

    def forward(self, x: torch.Tensor, y: torch.Tensor):
        feat_x, feat_y = self.net(x), self.net(y)
        feat_x = F.normalize(feat_x, dim=1)
        feat_y = F.normalize(feat_y, dim=1)

        return self.loss(feat_x, feat_y)